"""
Case Type   : 向量数据库可靠性  
Case Name   : hnsw索引在重建过程注入备机stop故障
Create At   : 2025/9
Owner       : l1azzzy
Description :
    1、建表
    2、插入大量数据
    3、创建索引
    4、步骤3执行过程中,stop备机
    5、恢复故障,重新执行
    6、主备使用索引查询
    7、主机重新构建索引
    8、主备均使用索引查询
    9、清理表
Expect      :
    1、建表成功
    2、插入数据成功
    3、创建索引成功
    4、stop备机成功
    5、故障恢复成功
    6、主备使用索引查询成功,主备查询结果相同
    7、主机重新构建索引成功
    8、主备均使用索引查询成功,查询结果相同
    9、清理表成功
History     :
"""

import os
import re
import time
import threading
import unittest

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

class DatavecIndexCase0029(unittest.TestCase):

    def setUp(self):
        self.log = Logger()
        self.log.info(f'----{os.path.basename(__file__)} start----')
        self.user_node = [CommonSH('PrimaryDbUser'), CommonSH('Standby1DbUser'), CommonSH('Standby2DbUser')]
        self.pri_user = Node(node='PrimaryDbUser')
        self.standby_user = Node(node='Standby1DbUser')
        self.com = Common('PrimaryDbUser')
        self.com_standby = Common('Standby1DbUser')
        self.constant = Constant()
        self.table_name = 't_datavec_index_case0029'
        self.index_name = 'idx_datavec_index_case0029'
        self.index_created = False
        self.log.info(f'----{os.path.basename(__file__)} end----')

    def stop_standby_db(self):
        """在后台执行stop备机数据库"""
        time.sleep(10)
        text = '-----step4:步骤3执行过程中,stop备机  expect:stop备机成功-----'
        self.log.info(text)
        stop_cmd = f'gs_ctl stop -D {macro.DB_INSTANCE_PATH} -m fast'
        res = self.standby_user.sh(stop_cmd)
        self.log.info(res)
        if 'server stopped' in res or 'waiting for server to shut down' in res:
            self.log.info("成功停止备机数据库")

    def test_datavec_index(self):
        text = '-----step1:建表  expect:建表成功-----'
        self.log.info(text)
        sql = f'create table {self.table_name} (id bigserial primary key, val vector(160));'
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, res)

        text = '-----step2:插入数据  expect:插入数据成功-----'
        self.log.info(text)
        sql = f'''
                    do $$
                    declare
                    vector_dim int := 160;
                    vector_value float[];
                    record_num int := 20000;
                    i int;
                    j int;
                    begin
                    for j in 1..record_num loop
                    for i in 1..vector_dim loop
                    vector_value[i] := ROUND(RANDOM() * 100, 2);
                    end loop;
                    insert into {self.table_name} (val) values (vector_value);
                    end loop;
                    end $$;
                    '''
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, res)

        text = '-----step3:创建索引  expect:创建索引成功-----'
        self.log.info(text)
        stop_thread = threading.Thread(target=self.stop_standby_db)
        stop_thread.daemon = True
        stop_thread.start()

        sql = f'''
        alter table {self.table_name} set(parallel_workers=30);
        set maintenance_work_mem ='1GB';
        create index {self.index_name} on {self.table_name} using hnsw (val vector_l2_ops) with(lists=1000);
        '''
        try:
            res = self.com.execut_db_sql(sql)
            self.log.info(res)
            if self.constant.INDEX_CREATE_SUCCESS in res:
                self.index_created = True
                self.log.info("索引创建在备机stop前已完成")
        except Exception as e:
            self.log.info(f"索引创建过程执行异常: {e}")

        stop_thread.join(timeout=10)

        text = '-----step5:恢复故障,重新执行  expect:故障恢复成功-----'
        self.log.info(text)
        
        restart_cmd = f'gs_ctl start -D {macro.DB_INSTANCE_PATH} -M standby'
        res = self.standby_user.sh(restart_cmd)
        self.log.info(res)
        self.assertTrue('server starting' in res or 'server started' in res or 'success' in res, "备机数据库启动失败")
        
        time.sleep(30)
        
        max_retries = 5
        for i in range(max_retries):
            try:
                check_sql = 'select 1;'
                res = self.com_standby.execut_db_sql(check_sql)
                if '1' in res:
                    self.log.info("备机数据库连接恢复成功")
                    sync_sql = 'select * from pg_stat_get_wal_senders();'
                    res_sync = self.com.execut_db_sql(sync_sql)
                    self.log.info(f"主备同步状态: {res_sync}")
                    
                    break
            except Exception as e:
                self.log.info(f"备机连接检查失败，重试 {i+1}/{max_retries}: {e}")
                time.sleep(10)
        else:
            self.fail("备机数据库连接恢复失败")
        if not self.index_created:
            sql = f'create index {self.index_name} on {self.table_name} using hnsw (val vector_l2_ops) with(lists=1000);'
            res = self.com.execut_db_sql(sql)
            self.log.info(res)
            self.assertIn(self.constant.INDEX_CREATE_SUCCESS, res)
            self.index_created = True
        else:
            self.log.info("索引已存在，无需重新创建")

        text = '-----step6:主备使用索引查询  expect:主备使用索引查询成功,主备查询结果相同-----'
        self.log.info(text)
        query_sql = f'''
        set enable_indexscan=on;
        set enable_seqscan=off;
        select id,val<->'[33.91,98.59,85.29,15.32,4.9,
        78.51,15.83,1.63,39.63,52,40.46,74.03,16.95,95.42,
        15.85,56.45,28.41,69.76,55.21,68.63,53.24,75.3,5.87,
        42.23,26.98,81.97,89.8,58.99,0.04,85.03,73.05,1.27,
        59.57,79.81,87.65,81.77,34,3.57,6.12,34.42,43.21,93.49,
        81.3,64.9,87.38,0.36,2.03,2.05,54.29,43.58,62.91,76.71,
        47.7,14.25,6.93,9.56,8.46,23.97,24.78,84.65,86.24,55.93,
        2.43,50,73.33,23.49,15.29,44.7,94.62,56.72,8.84,36.69,
        50.54,69.76,62.04,41.83,84.16,64.97,53.27,22.08,22.78,
        83.52,85.29,54.41,68.86,47.64,42.35,65.58,59.54,64.85,61.89,
        21.02,9.73,89.1,87.04,82.73,15.58,93.47,3.27,46.71,45.06,91.6,
        49.84,80.88,35.7,14.65,78.72,97.65,70,73.04,23.79,62.36,0.7,99.44,
        50.06,98.07,56.67,25.9,6.77,55.16,10.47,13.38,82.87,61.39,25.64,
        48.81,4.9,13.26,2.41,1.46,65.29,58.76,89.64,86.72,84.13,13.96,
        45.98,98.53,95.92,52.29,90.25,19.42,90.89,68.13,37.44,4.48,
        99.26,98.01,40.22,18.98,78.94,50.45,42.93,21.05,89.19,53.91,
        65.14,56.68,12.69,51.74]' from {self.table_name} 
        order by val<->'[33.91,98.59,85.29,15.32,4.9,78.51,15.83,1.63,
        39.63,52,40.46,74.03,16.95,95.42,15.85,56.45,28.41,69.76,55.21,
        68.63,53.24,75.3,5.87,42.23,26.98,81.97,89.8,58.99,0.04,85.03,
        73.05,1.27,59.57,79.81,87.65,81.77,34,3.57,6.12,34.42,43.21,
        93.49,81.3,64.9,87.38,0.36,2.03,2.05,54.29,43.58,62.91,76.71,
        47.7,14.25,6.93,9.56,8.46,23.97,24.78,84.65,86.24,55.93,2.43,50,
        73.33,23.49,15.29,44.7,94.62,56.72,8.84,36.69,50.54,69.76,62.04,
        41.83,84.16,64.97,53.27,22.08,22.78,83.52,85.29,54.41,68.86,47.64,
        42.35,65.58,59.54,64.85,61.89,21.02,9.73,89.1,87.04,82.73,15.58,
        93.47,3.27,46.71,45.06,91.6,49.84,80.88,35.7,14.65,78.72,
        97.65,70,73.04,23.79,62.36,0.7,99.44,50.06,98.07,56.67,25.9,
        6.77,55.16,10.47,13.38,82.87,61.39,25.64,48.81,4.9,13.26,2.41,
        1.46,65.29,58.76,89.64,86.72,84.13,13.96,45.98,98.53,95.92,52.29,
        0.25,19.42,90.89,68.13,37.44,4.48,99.26,98.01,40.22,18.98,78.94,
        50.45,42.93,21.05,89.19,53.91,65.14,56.68,12.69,51.74]' limit 10;
        '''
        res_primary = self.com.execut_db_sql(query_sql)
        self.log.info(f"主机查询结果: {res_primary}")
        self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_primary)
        res_standby = self.com_standby.execut_db_sql(query_sql)
        self.log.info(f"备机查询结果: {res_standby}")
        self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_standby)
        primary_results = self.extract_query_results(res_primary)
        standby_results = self.extract_query_results(res_standby)
        self.assertEqual(primary_results, standby_results, "主备查询结果不一致")

        text = '-----step7:主机重新构建索引  expect:主机重新构建索引成功-----'
        self.log.info(text)
        sql = f'reindex index {self.index_name};'
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.INDEX_CREATE_SUCCESS, res)

        text = '-----step8:主备均使用索引查询  expect:主备均使用索引查询成功,查询结果相同-----'
        self.log.info(text)
        res_primary_after_reindex = self.com.execut_db_sql(query_sql)
        self.log.info(f"重建索引后主机查询结果: {res_primary_after_reindex}")
        self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_primary_after_reindex)
        res_standby_after_reindex = self.com_standby.execut_db_sql(query_sql)
        self.log.info(f"重建索引后备机查询结果: {res_standby_after_reindex}")
        self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_standby_after_reindex)
        primary_results_after = self.extract_query_results(res_primary_after_reindex)
        standby_results_after = self.extract_query_results(res_standby_after_reindex)
        self.assertEqual(primary_results_after, standby_results_after, "重建索引后主备查询结果不一致")

    def extract_query_results(self, query_result):
        """从查询结果中提取id和距离值用于比较"""
        lines = query_result.split('\n')
        results = []
        for line in lines:
            if '|' in line and not line.startswith('---') and not line.startswith(' id ') and line.strip():
                parts = line.split('|')
                if len(parts) >= 2:
                    try:
                        id_val = parts[0].strip()
                        distance = parts[1].strip()
                        if id_val.isdigit():
                            results.append((int(id_val), float(distance)))
                    except (ValueError, IndexError):
                        continue
        return sorted(results)

    def tearDown(self):
        text = '-----step9:清理表  expect:清理表成功-----'
        self.log.info(text)
        sql = f'drop table if exists {self.table_name};'
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.DROP_TABLE_SUCCESS_MSG, res)
        self.log.info(f'----{os.path.basename(__file__)} end----')